#include <atomic>
#include <chrono>
#include <cstdint>
#include <thread>
#include <algorithm>

class TokenBucket {
public:
    using clock = std::chrono::steady_clock;
    using ns    = std::chrono::nanoseconds;

    // rate_per_sec: 每秒产⽣多少令牌（>=1）
    // capacity: 桶容量（可累积的最⼤令牌数，>=1）
    TokenBucket(double rate_per_sec, std::uint64_t capacity)
        : ns_per_token_(static_cast<int64_t>(1e9 / rate_per_sec)),
          burst_ns_(static_cast<int64_t>(capacity) * ns_per_token_),
          tat_ns_(to_ns(clock::now())) {
        if (rate_per_sec <= 0.0) throw std::invalid_argument("rate_per_sec must be > 0");
        if (capacity == 0)       throw std::invalid_argument("capacity must be >= 1");
        // 若 rate 不是 1e9 的整数因子，ns_per_token_ 是向下取整。
        // 如果你需要更高精度，可改用有理数/128-bit 累加；这里取整足够多数场景。
    }

    // 非阻塞获取 n 个令牌；成功返回 true，失败返回 false（不等待）
    bool try_acquire(std::uint64_t n = 1) {
        if (n == 0) return true;
        const int64_t need = static_cast<int64_t>(n) * ns_per_token_;
        auto now = to_ns(clock::now());

        int64_t old = tat_ns_.load(std::memory_order_relaxed);
        for (;;) {
            const int64_t base = std::max(old, now);
            const int64_t new_tat = base + need;

            // 等价令牌桶判断：不超过可累积的最⼤突发
            if (new_tat - now > burst_ns_) {
                return false; // 超出可用令牌（或突发上限）
            }

            // CAS 更新 TAT
            if (tat_ns_.compare_exchange_weak(
                    old, new_tat,
                    std::memory_order_acq_rel,
                    std::memory_order_relaxed)) {
                return true;
            }
            // CAS 失败时 old 已被刷新，继续尝试
        }
    }

    // 阻塞式获取 n 个令牌；返回实际等待的时间
    // 仍然全程无锁：可能短暂 sleep，未使用互斥量
    std::chrono::nanoseconds acquire(std::uint64_t n = 1) {
        if (n == 0) return ns::zero();
        const int64_t need = static_cast<int64_t>(n) * ns_per_token_;

        for (;;) {
            auto now = to_ns(clock::now());
            int64_t old = tat_ns_.load(std::memory_order_relaxed);
            const int64_t base = std::max(old, now);
            const int64_t new_tat = base + need;

            // 需要等待的时间 = max(0, new_tat - burst_ns - now)
            // 超过突发窗口的部分才需要等
            int64_t wait_ns = new_tat - burst_ns_ - now;
            if (wait_ns <= 0) wait_ns = 0;

            if (tat_ns_.compare_exchange_weak(
                    old, new_tat, std::memory_order_acq_rel, std::memory_order_relaxed)) {
                if (wait_ns > 0) {
                    std::this_thread::sleep_for(ns(wait_ns));
                }
                return ns(wait_ns);
            }
            // CAS 失败则重试
        }
    }

    // 估算当前可用令牌（仅近似，供监控/指标使用）
    double approx_tokens() const {
        auto now = to_ns(clock::now());
        int64_t tat = tat_ns_.load(std::memory_order_relaxed);
        int64_t lag = now - tat; // <0 表示“欠账”；>0 表示“提前/积累”
        // 将 "提前" 时间转换成“可用令牌数”
        double tokens = (double)lag / (double)ns_per_token_ + capacity_tokens();
        if (tokens < 0) tokens = 0;
        double cap = (double)capacity_tokens();
        if (tokens > cap) tokens = cap;
        return tokens;
    }

    // 配置查询
    double rate_per_sec() const { return 1e9 / (double)ns_per_token_; }
    std::uint64_t capacity_tokens() const { return (std::uint64_t)(burst_ns_ / ns_per_token_); }

private:
    static int64_t to_ns(clock::time_point t) {
        return std::chrono::duration_cast<ns>(t.time_since_epoch()).count();
    }

    const int64_t ns_per_token_;   // 每个令牌对应的纳秒
    const int64_t burst_ns_;       // 可累积的最⼤突发（纳秒）
    std::atomic<int64_t> tat_ns_;  // 理论到达时间 TAT（纳秒时间戳）
};

// ======================= 示例 =======================
// 多线程并发限流：总速率 10 rps，容量 20 令牌。
#include <iostream>
#include <vector>

int main() {
    TokenBucket limiter(10.0, 20); // 10个/秒，最多积累20个
    constexpr int kThreads = 4;
    constexpr int kOpsPerThread = 25;

    std::vector<std::thread> ths;
    for (int i = 0; i < kThreads; ++i) {
        ths.emplace_back([&, i] {
            for (int j = 0; j < kOpsPerThread; ++j) {
                if (limiter.try_acquire()) {
                    // 成功拿到 1 个令牌
                    // … 执行业务 …
                } else {
                    // 没拿到就等一下再试（也可以直接调用阻塞式 acquire）
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                    j--; // 重试本次
                }
            }
            std::cout << "Thread " << i << " done\n";
        });
    }
    for (auto &t : ths) t.join();
    std::cout << "All done\n";
    return 0;
}
